import os
import numpy as np
import tensorflow as tf
import mymodel
import input_data
import alexnet
from tensorflow.python.framework.graph_util import convert_variables_to_constants
# %%

N_CLASSES = 2
IMG_W = 200  # resize the image, if the input image is too large, training will be very slow.
IMG_H = 200
BATCH_SIZE = 130
CAPACITY = 2000
MAX_STEP = 20000  # with current parameters, it is suggested to use MAX_STEP>10k
learning_rate = 0.001  # with current parameters, it is suggested to use learning rate<0.0001


# %%
def run_training():
    # you need to change the directories to yours.
    train_dir = 'D:/scratch/train3/'
    logs_train_dir = 'E:/Projects/tensorflow/python/logs/train_learningrate1/'

    train,train_label = input_data.get_files(train_dir)

    train_batch,train_label_batch = input_data.get_batch(train,
    train_label,
    IMG_W,
    IMG_H,
    BATCH_SIZE,
    CAPACITY)
    input_images = tf.reshape(train_batch, shape=[BATCH_SIZE, IMG_W, IMG_H, 1],name="input")
    train_logits,pool3 = alexnet.inference(input_images,BATCH_SIZE,N_CLASSES,0.5)
    train_loss = mymodel.losses(train_logits,train_label_batch)
    train_op = mymodel.trainning(train_loss,learning_rate)
    train__acc = mymodel.evaluation(train_logits,train_label_batch)

    summary_op = tf.summary.merge_all()
    sess = tf.Session()
    train_writer = tf.summary.FileWriter(logs_train_dir, sess.graph)
    saver = tf.train.Saver()

    sess.run(tf.global_variables_initializer())
    coord = tf.train.Coordinator()
    threads = tf.train.start_queue_runners(sess=sess, coord=coord)

    try:
        # print(pool3)
        # for step in np.arange(MAX_STEP):
        #     if coord.should_stop():
        #         break
        #     _, tra_loss, tra_acc = sess.run([train_op, train_loss, train__acc])
        #     if step % 50 == 0:
        #         print('Step %d, train loss = %.2f, train accuracy = %.2f%%' % (step, tra_loss, tra_acc * 100.0))
        #         summary_str = sess.run(summary_op)
        #         train_writer.add_summary(summary_str, step)
        #
        #     if step % 2000 == 0 or (step + 1) == MAX_STEP:
        #         checkpoint_path = os.path.join(logs_train_dir, 'model.ckpt')
        #         saver.save(sess, checkpoint_path, global_step=step)
        tf.train.write_graph(sess.graph, logs_train_dir, 'expert_graph1.pb', as_text=False)
        graph = convert_variables_to_constants(sess, sess.graph_def, ["output/output"])
        tf.train.write_graph(graph, logs_train_dir, 'expert-graph.pb', as_text=False)
    except tf.errors.OutOfRangeError:
        print('Done training -- epoch limit reached')
    finally:
        coord.request_stop()

    coord.join(threads)
    sess.close()


    #%% Evaluate one image
    #when training, comment the following codes.


from PIL import Image
import matplotlib.pyplot as plt

def get_one_image(train):
       '''Randomly pick one image from training data
       Return: ndarray
       '''
       n = len(train)
       ind = np.random.randint(0, n)
       img_dir = train[ind]

       image = Image.open(img_dir)
       plt.imshow(image)
       plt.show()
       image = image.resize([200, 200])
       image = np.array(image)
       return image


run_training()

